{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8843ef92",
   "metadata": {},
   "source": [
    "\n",
    "# 课程大纲\n",
    "## 大语言模型直接生成种子实验\n",
    "## 通过种子生成器生成种子实验\n",
    "## 种子生成自动化流程实验\n",
    "## 种子生成器规模化流程实验\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39eeba6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 基本环境\n",
    "\n",
    "from openai import OpenAI\n",
    "key = \"\" #换成你的key\n",
    "url = \"\" # 换成你的url\n",
    "model_name = \"\" # 换成你的模型名，建议使用qwen3-32b以上的模型\n",
    "# 大语言模型回显\n",
    "def print_llm(text):\n",
    "    print(text, end=\"\", flush=True)\n",
    "\n",
    "def run_llm(prompt_str):\n",
    "    result = \"\"\n",
    "    client = OpenAI(api_key=key, base_url=url)\n",
    "    response = client.chat.completions.create(\n",
    "        model=model_name,\n",
    "                    messages=[\n",
    "                        {\"role\": \"user\", \"content\": prompt_str},\n",
    "                    ],\n",
    "                    stream=True,\n",
    "                    # 把timeout值设置高一些\n",
    "                    timeout=600\n",
    "                )\n",
    "    for chunk in response:\n",
    "        content = chunk.choices[0].delta.content\n",
    "        if content is not None:\n",
    "            result += content\n",
    "            print_llm(content)\n",
    "    return result\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61fab620",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试大语言模型接口是否工作正常\n",
    "result = run_llm(\"你好\\\\no_think\")\n",
    "print(\"=\"*40)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19cccade",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chapter 1 直接生成jxl字节流的测试用例\n",
    "prompt_str = \"\"\"请以十六进制转义字符串(\"\\\\x00\\\\x01\")的格式生成10个jxl格式的字节流。\n",
    "要求尽可能覆盖到jxl解析代码深层的路径。\n",
    "要求：\n",
    "输出的必须是\"\\\\x00\\\\x01\"这种十六进制转义字符串。\n",
    "输出的字符串要包含jxl文件所必须的各种magic number，以及符合jxl文件的格式。\n",
    "生成的json格式不要包含注释。\n",
    "Think step by step and output your thinking progress.\n",
    "最终以json格式输出，下面是示例：\n",
    "```json\n",
    "{\n",
    "    \"jxl_bytes\": [十六进制转义字符串列表]    \n",
    "}\n",
    "\\\\no_think\n",
    "```\n",
    "\"\"\"\n",
    "result = run_llm(prompt_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49392366",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chapter 1 直接生成jxl字节的测试用例\n",
    "import os\n",
    "# 加入解析python的代码\n",
    "def extract_json(text):\n",
    "    import re\n",
    "    import json\n",
    "    pattern = r'```json\\n([\\s\\S]*?)\\n```'\n",
    "    matches = re.findall(pattern, text, flags=re.DOTALL)\n",
    "    valid_json = []\n",
    "    for match in matches:\n",
    "        try:\n",
    "            # 尝试解析为JSON\n",
    "            json_obj = json.loads(match)\n",
    "            valid_json.append(json_obj)\n",
    "        except json.JSONDecodeError:\n",
    "            # 无效JSON跳过\n",
    "            continue\n",
    "    return valid_json\n",
    "\n",
    "seed_dict = extract_json(result)[-1]\n",
    "seeds = seed_dict['jxl_bytes']\n",
    "target_floder = \"./j40/corpus\" # 换成你的目标代码位置\n",
    "for i,seed in enumerate(seeds):\n",
    "    file_name = os.path.join(target_floder,f\"seed_{i}.jxl\")\n",
    "    with open(file_name,'wb') as f:\n",
    "        f.write(bytes(seed,encoding='utf-8'))\n",
    "    print(bytes(seed,encoding='utf-8'))\n",
    "\n",
    "## 代码生成后，需要运行libfuzzer\n",
    "## cd j40\n",
    "## mkdir corpus\n",
    "## ./test ./corpus -runs=2 -detect_leaks=0\n",
    "## 查看覆盖率\n",
    "## llvm-profdata merge -sparse *.profraw -o default.profdata\n",
    "## llvm-cov report test -instr-profile=default.profdata\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c433e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chapter1 将流程合并\n",
    "prompt_str = \"\"\"请以十六进制转义字符串(\"\\\\x00\\\\x01\")的格式生成10个jxl格式的字节流。\n",
    "要求尽可能覆盖到jxl解析代码深层的路径。\n",
    "要求：\n",
    "输出的必须是\"\\\\x00\\\\x01\"这种十六进制转义字符串。\n",
    "输出的字符串要包含jxl文件所必须的各种magic number，以及符合jxl文件的格式。\n",
    "生成的json格式不要包含注释。\n",
    "Think step by step and output your thinking progress.\n",
    "最终以json格式输出，下面是示例：\n",
    "```json\n",
    "{\n",
    "    \"jxl_bytes\": [十六进制转义字符串列表]    \n",
    "}\n",
    "\\\\no_think\n",
    "```\n",
    "\"\"\"\n",
    "max_count = 20\n",
    "count = 0\n",
    "target_floder = \"./j40/corpus\" # 换成你的目标代码位置\n",
    "while True:\n",
    "    result = run_llm(prompt_str)\n",
    "    try:\n",
    "        seed_dict = extract_json(result)[-1]\n",
    "        seeds = seed_dict['jxl_bytes']\n",
    "        for seed in seeds:\n",
    "            if count > max_count:\n",
    "                break\n",
    "            file_name = os.path.join(target_floder,f\"seed_{count}.jxl\")\n",
    "            with open(file_name,'wb') as f:\n",
    "                f.write(bytes(seed,encoding='utf-8'))\n",
    "            count += 1\n",
    "            print(bytes(seed,encoding='utf-8'))\n",
    "    except:\n",
    "        pass\n",
    "print(\"生成成功\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b59bf31b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chapter 2 生成种子生成器的方式产生种子\n",
    "\n",
    "## chapter 1中，发现种子生成效果不佳，还容易让LLM陷入循环输出的bug\n",
    "## 因此采用生成种子生成器的方法，这部分模拟G2Fuzz\n",
    "\n",
    "## Step 1 获取jxl特性\n",
    "prompt_str = \"\"\"根据你对jxl图像知识，列出jxl图像在文件格式上的10种算法特性。\n",
    "\\\\no_think\n",
    "\"\"\"\n",
    "feature_result = run_llm(prompt_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "147b1348",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chapter 2 \n",
    "## Step 2 直接根据特性生成python脚本\n",
    "prompt_str = f\"\"\"根据下面的jxl格式要求，写一个Python脚本，用来生成满足下面要求的jxl文件,要求要使用pillow_jxl库。\n",
    "要求：1.输出的文件夹要用户指定。\n",
    "2. 输出的函数接口名为run,里面要有输出文件夹、输出文件数量参数两个参数。\n",
    "3. 整个程序只有一个run函数，所有生成逻辑都要放到run函数中。\n",
    "4. 除了pillow_jxl库外，只能使用struct，numpy,os,sys库，严禁使用其他库。\n",
    "5. 不需要加入日志或者其他与生成jxl文件无关的功能。\n",
    "6. 不要加入任何调用run函数的示例代码。\n",
    "\n",
    "{feature_result}\n",
    "生成的代码是python代码：\n",
    "```python\n",
    "//代码放到这里\n",
    "```\n",
    "\\\\no_think\n",
    "\"\"\"\n",
    "result = run_llm(prompt_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e51f7c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chapter 2 \n",
    "## Step 3 手工 测试python代码\n",
    "## 测试python生成代码\n",
    "## 【NOTE】 把上一步的代码手动复制到这里\n",
    "import os\n",
    "import sys\n",
    "import struct\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import pillow_jxl\n",
    "\n",
    "def run(output_folder, num_files):\n",
    "    if not os.path.exists(output_folder):\n",
    "        os.makedirs(output_folder)\n",
    "\n",
    "    # 生成简单的测试图像\n",
    "    width, height = 256, 256\n",
    "    for i in range(num_files):\n",
    "        # 创建一个简单的渐变图像\n",
    "        gradient = np.zeros((height, width, 3), dtype=np.uint8)\n",
    "        for y in range(height):\n",
    "            for x in range(width):\n",
    "                gradient[y, x] = [x % 256, y % 256, (x + y) % 256]\n",
    "\n",
    "        # 使用Pillow创建图像\n",
    "        img = Image.fromarray(gradient, 'RGB')\n",
    "\n",
    "        # 保存为JXL格式\n",
    "        output_path = os.path.join(output_folder, f\"image_{i:04d}.jxl\")\n",
    "        img.save(output_path, format='JXL')\n",
    "\n",
    "\n",
    "target_floder = \"./j40/corpus\"\n",
    "run(target_floder, 40)\n",
    "\n",
    "## 然后再按照chapter 1最后的步骤获取覆盖率"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d61ae120",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chapter3 种子生成器流程自动化\n",
    "\n",
    "def extract_python_code(prompt_str):\n",
    "    import re\n",
    "    python_code = re.findall(r\"```python(.*?)```\", result, re.DOTALL)[0]\n",
    "    return python_code\n",
    "# 固定的框架代码，用于调用run函数\n",
    "wrapper_code = \"\"\"\n",
    "\n",
    "target_floder = \"./j40/corpus\"\n",
    "# 生成数量设为40\n",
    "gen_num = 40\n",
    "run(target_floder, gen_num)\n",
    "\"\"\"\n",
    "\n",
    "# 生成特性\n",
    "feature_prompt_str = \"\"\"根据你对jxl图像知识，列出jxl图像在文件格式上的10种算法特性。\n",
    "\\\\no_think\n",
    "\"\"\"\n",
    "feature_result = run_llm(feature_prompt_str)\n",
    "\n",
    "# 生成python代码\n",
    "\n",
    "code_gen_prompt_str = f\"\"\"根据下面的jxl格式要求，写一个Python脚本，用来生成满足下面要求的jxl文件,要求要使用pillow_jxl库。\n",
    "要求：1.输出的文件夹要用户指定。\n",
    "2. 输出的函数接口名为run,里面要有输出文件夹、输出文件数量参数两个参数。\n",
    "3. 整个程序只有一个run函数，所有生成逻辑都要放到run函数中。\n",
    "4. 除了pillow_jxl库外，只能使用struct，numpy,os,sys库，严禁使用其他库。\n",
    "5. 不需要加入日志或者其他与生成jxl文件无关的功能。\n",
    "6. 不要加入任何调用run函数的示例代码。\n",
    "\n",
    "{feature_result}\n",
    "生成的代码是python代码：\n",
    "```python\n",
    "//代码放到这里\n",
    "```\n",
    "\\\\no_think\n",
    "\"\"\"\n",
    "result = run_llm(code_gen_prompt_str)\n",
    "python_code = extract_python_code(result)\n",
    "\n",
    "print(\"\\n\"*5)\n",
    "print(python_code)\n",
    "\n",
    "# 自动化流程，运行代码\n",
    "run_code = python_code + wrapper_code\n",
    "exec(run_code)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e294335b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chapter 4 加入校验和修改\n",
    "\n",
    "# chapter 3中，由于没有对代码的校验机制，导致有时候会执行失败。\n",
    "# 本章节加入校验和修改机制，提升种子生成成功概率。\n",
    "\n",
    "# 校验代码有没有生成\n",
    "# 校验生成的代码语法是否有问题\n",
    "# 校验生成的代码能否正常运行\n",
    "# 校验生成的代码能够正常产生种子\n",
    "\n",
    "import ast\n",
    "wrapper_code = \"\"\"\n",
    "\n",
    "\n",
    "target_floder = \"./j40/corpus\"\n",
    "# 生成数量设为40\n",
    "gen_num = 40\n",
    "run(target_floder, gen_num)\n",
    "\"\"\"\n",
    "def extract_python_code(prompt_str):\n",
    "    import re\n",
    "    python_code = re.findall(r\"```python(.*?)```\", prompt_str, re.DOTALL)[0]\n",
    "    return python_code\n",
    "# step1 将生成部分包装成函数\n",
    "def gen_seed_generator_code(feature_result):\n",
    "    # 生成python代码\n",
    "\n",
    "    code_gen_prompt_str = f\"\"\"根据下面的jxl格式要求，写一个Python脚本，用来生成满足下面要求的jxl文件,要求要使用pillow_jxl库。\n",
    "    要求：1.输出的文件夹要用户指定。\n",
    "    2. 输出的函数接口名为run,里面要有输出文件夹、输出文件数量参数两个参数。\n",
    "    3. 整个程序只有一个run函数，所有生成逻辑都要放到run函数中。\n",
    "    4. 除了pillow_jxl库外，只能使用struct，numpy,os,sys库，严禁使用其他库。\n",
    "    5. 不需要加入日志或者其他与生成jxl文件无关的功能。\n",
    "    6. 不要加入任何调用run函数的示例代码。\n",
    "\n",
    "    {feature_result}\n",
    "    生成的代码是python代码：\n",
    "    ```python\n",
    "    //代码放到这里\n",
    "    ```\n",
    "    \\\\no_think\n",
    "    \"\"\"\n",
    "    result = run_llm(code_gen_prompt_str)\n",
    "    return result\n",
    "def generate_featrue():\n",
    "    # 生成特性\n",
    "    feature_prompt_str = \"\"\"根据你对jxl图像知识，列出jxl图像在文件格式上的10种算法特性。\n",
    "    \\\\no_think\n",
    "    \"\"\"\n",
    "    feature_result = run_llm(feature_prompt_str)\n",
    "    return feature_result\n",
    "    \n",
    "# 先生成特性\n",
    "feature_result = generate_featrue()\n",
    "\n",
    "# step 2 设立死循环，不断尝试生成代码，直到成功为止\n",
    "target_floder = \"./j40/corpus\"\n",
    "while True:\n",
    "    rough_code = gen_seed_generator_code(feature_result)\n",
    "    python_code = extract_python_code(rough_code)\n",
    "    # 如果无法解析到python代码，则重新生成\n",
    "    if not python_code: continue\n",
    "    run_code = python_code + wrapper_code\n",
    "    print(run_code)\n",
    "    # 通过python ast模块的parse函数，判断生成的代码是否有语法问题\n",
    "    try:\n",
    "        ast.parse(run_code)\n",
    "    except:\n",
    "        print(\"\\nInvalid Python Code\\n\")\n",
    "        continue\n",
    "    # 尝试运行目标代码，如果成功，则检查是否生成了种子\n",
    "    try:\n",
    "        exec(run_code)\n",
    "        file_list = os.listdir(target_floder)\n",
    "        if len(file_list) > 0:\n",
    "            print(\"\\n种子生成成功\\n\")\n",
    "            break\n",
    "        else:\n",
    "            print(\"\\n生成代码没有产生种子\\n\")\n",
    "    except:\n",
    "        print(\"\\n生成代码运行错误\\n\")\n",
    "        continue\n",
    "\n",
    "## 代码生成成功后通过chapter 1内容运行代码并获取覆盖率。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4efd518a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# chapter 5 加入对脚本的选择\n",
    "# chapter4中，在能够生成种子后就停止，实际上，还是需要对生成种子的质量筛选，选择效果最好的种子生成器。\n",
    "# 产生多个生成器，筛选效果最好的一个作为最终结果\n",
    "import ast\n",
    "import sys\n",
    "import os\n",
    "wrapper_code = \"\"\"\n",
    "\n",
    "\n",
    "target_floder = \"./j40/corpus\"\n",
    "# 生成数量设为40\n",
    "gen_num = 40\n",
    "run(target_floder, gen_num)\n",
    "\"\"\"\n",
    "# 更改回显方式\n",
    "global print_count\n",
    "print_count = 0\n",
    "def print_llm(text):\n",
    "    global print_count\n",
    "    symbols = ['|','|','|','|','|', '/', '/', '/', '/', '/', '-','-','-','-','-', '\\\\','\\\\','\\\\','\\\\','\\\\']\n",
    "    print_count += 1\n",
    "    print_count = print_count % len(symbols)\n",
    "    sys.stdout.write('\\r大语言模型工作中 ' + symbols[print_count])\n",
    "    sys.stdout.flush()\n",
    "    \n",
    "    \n",
    "def extract_python_code(prompt_str):\n",
    "    import re\n",
    "    python_code = re.findall(r\"```python(.*?)```\", prompt_str, re.DOTALL)[0]\n",
    "    return python_code\n",
    "# step1 将生成部分包装成函数\n",
    "def gen_seed_generator_code(feature_result):\n",
    "    # 生成python代码\n",
    "\n",
    "    code_gen_prompt_str = f\"\"\"根据下面的jxl格式要求，写一个Python脚本，用来生成满足下面要求的jxl文件,要求要使用pillow_jxl库。\n",
    "    要求：1.输出的文件夹要用户指定。\n",
    "    2. 输出的函数接口名为run,里面要有输出文件夹、输出文件数量参数两个参数。\n",
    "    3. 整个程序只有一个run函数，所有生成逻辑都要放到run函数中。\n",
    "    4. 除了pillow_jxl库外，只能使用struct，numpy,os,sys库，严禁使用其他库。\n",
    "    5. 不需要加入日志或者其他与生成jxl文件无关的功能。\n",
    "    6. 不要加入任何调用run函数的示例代码。\n",
    "\n",
    "    {feature_result}\n",
    "    生成的代码是python代码：\n",
    "    ```python\n",
    "    //代码放到这里\n",
    "    ```\n",
    "    \\\\no_think\n",
    "    \"\"\"\n",
    "    result = run_llm(code_gen_prompt_str)\n",
    "    return result\n",
    "def generate_featrue():\n",
    "    # 生成特性\n",
    "    feature_prompt_str = \"\"\"根据你对jxl图像知识，列出jxl图像在文件格式上的10种算法特性。\n",
    "    \\\\no_think\n",
    "    \"\"\"\n",
    "    feature_result = run_llm(feature_prompt_str)\n",
    "    return feature_result\n",
    "\n",
    "\n",
    "# 获取覆盖率\n",
    "def get_cov():\n",
    "    run_target_cmd = \"./j40/test ./j40/corpus -runs=2 -detect_leaks=0\"\n",
    "    os.system(run_target_cmd)\n",
    "    os.system(\"llvm-profdata merge -sparse *.profraw -o default.profdata\")\n",
    "    get_cov_cmd = \"llvm-cov report ./j40/test -instr-profile=default.profdata > cov_info\"\n",
    "    os.system(get_cov_cmd)\n",
    "    with open(\"cov_info\", \"r\") as f:\n",
    "        cov_info = f.read()\n",
    "    cov_line = cov_info.split(\"TOTAL\")[1]\n",
    "    cov_list = cov_line.split(\" \")\n",
    "    cov = cov_list[-1].split(\"%\")[0]\n",
    "    os.system(\"rm default.prof*\")\n",
    "    os.system(\"rm cov_info\")\n",
    "    return float(cov)\n",
    "\n",
    "\n",
    "# 先生成特性\n",
    "feature_result = generate_featrue()\n",
    "\n",
    "# step 2 设立死循环，不断尝试生成代码，直到成功为止\n",
    "target_floder = \"./j40/corpus/*\"\n",
    "rm_cmd =f\"rm {target_floder}\"\n",
    "\n",
    "# 设置全局变量\n",
    "best_python_code = \"\"\n",
    "success_gen_count = 0\n",
    "MAX_GEN_COUNT = 10\n",
    "MAX_COV = 0\n",
    "while True:\n",
    "    os.system(rm_cmd)\n",
    "    rough_code = gen_seed_generator_code(feature_result)\n",
    "    python_code = extract_python_code(rough_code)\n",
    "    # 如果无法解析到python代码，则重新生成\n",
    "    if not python_code: continue\n",
    "    run_code = python_code + wrapper_code\n",
    "    # print(run_code)\n",
    "    # 解析目标语法是否有问题\n",
    "    try:\n",
    "        ast.parse(run_code)\n",
    "    except:\n",
    "        print(\"\\n生成代码语法错误\\n\")\n",
    "        continue\n",
    "    # 尝试运行目标代码，如果成功，则检查是否生成了种子\n",
    "    try:\n",
    "        exec(run_code)\n",
    "        file_list = os.listdir(target_floder)\n",
    "        if len(file_list) > 0:\n",
    "            print(\"\\n种子生成成功\\n\")\n",
    "            success_gen_count += 1\n",
    "            cov = get_cov()\n",
    "            print(f\"当前覆盖率: {cov} and 最大覆盖率: {MAX_COV}\")\n",
    "            if cov > MAX_COV:\n",
    "                best_python_code = run_code\n",
    "                MAX_COV = cov\n",
    "            if success_gen_count > MAX_GEN_COUNT:\n",
    "                break\n",
    "        else:\n",
    "            print(\"\\n生成代码没有产生种子\\n\")\n",
    "    except:\n",
    "        print(\"\\n生成代码运行错误\\n\")\n",
    "        continue"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "infosec",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
